import os
import json
from django.core.management.base import BaseCommand
from api.models import Server

class Command(BaseCommand):
    help = '从JSON文件导入MCP Server 数据'

    def add_arguments(self, parser):
        parser.add_argument('--file', type=str, default='../mcphub-scripts/mcp-server-list-tools/mcp-list-v2.json', help='JSON文件路径')
        parser.add_argument('--fix-category', action='store_true', help='修复category字段')

    def handle(self, *args, **options):
        # 如果需要修复category字段
        if options.get('fix_category'):
            return self.fix_category_field()
            
        file_path = options['file']
        
        if not os.path.exists(file_path):
            self.stdout.write(self.style.ERROR(f'文件不存在: {file_path}'))
            return
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                servers_data = json.load(f)
        except Exception as e:
            self.stdout.write(self.style.ERROR(f'读取JSON文件出错: {str(e)}'))
            return
        
        # 计数器
        created_count = 0
        updated_count = 0
        error_count = 0
        
        for server_data in servers_data:
            try:
                server_id = server_data.get('id')
                if not server_id:
                    continue

                # 设置必要的默认值
                if not server_data.get('tags'):
                    server_data['tags'] = []
                
                # 创建或更新 Server 
                server, created = Server.objects.update_or_create(
                    id=server_id,
                    defaults={
                        'name': server_data.get('name', ''),
                        'name_cn': server_data.get('name_cn', ''),
                        'icon': server_data.get('icon', ''),
                        'provider': server_data.get('provider', ''),
                        'description': server_data.get('description', ''),
                        'description_cn': server_data.get('description_cn', ''),
                        'category': server_data.get('category', []),
                        'category_cn': server_data.get('category_cn', []),
                        'tags': server_data.get('tags', []),
                        'is_official': server_data.get('is_official', False),
                        'is_internal': server_data.get('is_internal', False),
                        'auth_method': server_data.get('auth_method', ''),
                        'dev_language': server_data.get('dev_language', []),
                        'from_site': server_data.get('from_site', ''),
                        'from_site_url': server_data.get('from_site_url', ''),
                        'server_config': server_data.get('server_config', ''),
                        'env_schema': server_data.get('env_schema', {}),
                        'tools': server_data.get('tools', []),
                        'readme': server_data.get('readme', ''),
                        'readme_cn': server_data.get('readme_cn', ''),
                        'verified': server_data.get('verified', False),  # 注意源数据中是'verifed'而不是'verified'
                        'zone': server_data.get('zone', ''),
                        'internal_doc_url': server_data.get('internal_doc_url', ''),
                        'is_deleted': server_data.get('is_deleted', False),
                    }
                )
                
                if created:
                    created_count += 1
                    self.stdout.write(self.style.SUCCESS(f'创建新 Server: {server.name}'))
                else:
                    updated_count += 1
                    self.stdout.write(self.style.SUCCESS(f'更新 Server: {server.name}'))
                    
            except Exception as e:
                error_count += 1
                self.stdout.write(self.style.ERROR(f'导入 Server {server_data.get("id", "未知")} 时出错: {str(e)}'))

        self.stdout.write(self.style.SUCCESS(f'导入完成: 创建 {created_count} 个, 更新 {updated_count} 个, 错误 {error_count} 个'))
        
    def fix_category_field(self):
        """修复数据库中已有记录的category字段，将字符串转换为JSON数组"""
        servers = Server.objects.all()
        fixed_count = 0
        error_count = 0
        
        for server in servers:
            try:
                # 检查category字段是否为字符串
                if isinstance(server.category, str):
                    # 将字符串转换为列表
                    category_name = server.category
                    server.category = [category_name]
                    server.save(update_fields=['category'])
                    fixed_count += 1
                    self.stdout.write(self.style.SUCCESS(f'修复 Server {server.name} 的category字段'))
            except Exception as e:
                error_count += 1
                self.stdout.write(self.style.ERROR(f'修复 Server {server.name} 的category字段时出错: {str(e)}'))
                
        self.stdout.write(self.style.SUCCESS(f'修复完成: 修复 {fixed_count} 个, 错误 {error_count} 个')) 